from neo4j import GraphDatabase
from neo4j.exceptions import Neo4jError
from typing import List, Any, Union, Dict
from Node import Clazz, Method, TestClazz, _convert_node
import threading
from LoggerManager import LoggerManager

# danger dependency
from Config import *


class CKGRetriever:
    """Singleton CKGRetriever that manages Neo4j connections"""
    _instance = None
    _lock = threading.Lock()  # Thread lock to ensure thread safety

    def __new__(cls, uri, user, password):
        with cls._lock:
            if cls._instance is None:
                cls._instance = super(CKGRetriever, cls).__new__(cls)
                cls._instance._init_driver(uri, user, password)
        return cls._instance

    def _init_driver(self, uri, user, password):
        """Initialize Neo4j driver"""
        self.driver = GraphDatabase.driver(uri, auth=(user, password))
        self.focal_method_id = -1

    def close(self):
        """Close Neo4j connection"""
        if self.driver:
            self.driver.close()

    def change_focal_method_id(self, focal_method_id):
        """
        Change the current focal_method_id.
        """
        self.focal_method_id = focal_method_id


    def lazy_load_classes(self, batch_size=10):
        """
        Generator to lazily load clazz nodes from the Neo4j database.

        :param batch_size: Number of nodes to fetch in each batch.
        :return: Yields clazz nodes one by one.
        """
        with self.driver.session() as session:
            offset = 0
            while True:
                # Run a Cypher query with LIMIT and SKIP for pagination
                result = session.run(
                    """
                    MATCH (c:Clazz) 
                    RETURN id(c) AS node_id, c
                    ORDER BY c.name ASC
                    SKIP $offset LIMIT $batch_size
                    """,
                    offset=offset, batch_size=batch_size
                )

                # Convert results to a list
                result = list(result)

                # If no more results, stop iteration
                if not result:
                    break

                # Yield each clazz one by one
                for record in result:
                    yield _convert_node(record["c"])

                # Move to the next batch
                offset += batch_size

    def load_methods(self, focal_method=True):
        """
        Generator to lazily load method nodes from the Neo4j database.

        :param focal_method: return only focal method if True which have public access modifier and not abstract
        :param batch_size: Number of nodes to fetch in each batch.
        :return: Yields method nodes one by one.
        """
        with self.driver.session() as session:
            # Run a Cypher query with LIMIT and SKIP for pagination
            #     WHERE m.modifiers CONTAINS "public" AND NOT m.modifiers CONTAINS "abstract" AND m.type = "METHOD"
            if focal_method:
                result = session.run(
                    """
                    MATCH (m:Method) 
                    WHERE m.modifiers CONTAINS "public" AND NOT m.modifiers CONTAINS "abstract" AND m.type = "METHOD"
                    RETURN id(m) AS node_id, m
                    ORDER BY m.name ASC
                    """
                )
                result = list(result)
            else:
                result = session.run(
                    """
                    MATCH (m:Method) 
                    RETURN id(m) AS node_id, m
                    ORDER BY m.name ASC
                    """
                )
                result = list(result)
            return [_convert_node(record["m"]) for record in result]
        
    def load_methods_ctx(self, fake_number=10):
        """
        Generator to lazily load method nodes from the Neo4j database.

        :param focal_method: return only focal method if True which have public access modifier and not abstract
        :param batch_size: Number of nodes to fetch in each batch.
        :return: Yields method nodes one by one.
        """
        with self.driver.session() as session:
            # Run a Cypher query with LIMIT and SKIP for pagination
            #     WHERE m.modifiers CONTAINS "public" AND NOT m.modifiers CONTAINS "abstract" AND m.type = "METHOD"
            result = session.run(
                """
                MATCH (m:Method) <-[:CONTAIN]- (c:Clazz)
                WHERE m.modifiers CONTAINS "public" AND NOT m.modifiers CONTAINS "abstract" AND m.type = "METHOD"
                RETURN id(m) AS node_id, m, c.content AS ctx
                ORDER BY m.name ASC
                """
            )
            result = list(result)
            
        return [_convert_node(record["m"], record["ctx"]) for record in result]

    def load_filtered_methods(self):
        """
        Load 100 pseudo-randomly selected method nodes from the Neo4j database,
        filtered by modifiers and sorted deterministically by a hash of the node ID.

        :return: A list of converted method nodes.
        """
        with self.driver.session() as session:
            result = session.run(
                """
                MATCH (m:Method)
                WHERE m.modifiers CONTAINS "public" 
                  AND NOT m.modifiers CONTAINS "abstract" 
                  AND m.type = "METHOD"
                RETURN id(m) AS node_id, m
                ORDER BY toInteger(abs(id(m) * 17 + 12345))  // deterministic pseudo-random sort
                LIMIT 200
                """
            )
            result = list(result)
            return [_convert_node(record["m"]) for record in result]

    def lazy_load_variables(self, batch_size=10):
        """
        Generator to lazily load variable nodes from the Neo4j database.

        :param batch_size: Number of nodes to fetch in each batch.
        :return: Yields variable nodes one by one.
        """
        with self.driver.session() as session:
            offset = 0
            while True:
                # Run a Cypher query with LIMIT and SKIP for pagination
                result = session.run(
                    """
                    MATCH (v:Variable) 
                    RETURN id(v) AS node_id, v
                    ORDER BY v.name ASC
                    SKIP $offset LIMIT $batch_size
                    """,
                    offset=offset, batch_size=batch_size
                )
                result = list(result)

                # Convert results to a list
                variables = list(result)

                # If no more results, stop iteration
                if not variables:
                    break

                # Yield each variable one by one
                for record in variables:
                    yield _convert_node(record["v"])

                # Move to the next batch
                offset += batch_size

    def lazy_load_test_classes(self, batch_size=10):
        """
        Generator to lazily load test class nodes from the Neo4j database.

        :param batch_size: Number of nodes to fetch in each batch.
        :return: Yields test class nodes one by one.
        """
        with self.driver.session() as session:
            offset = 0
            while True:
                # Run a Cypher query with LIMIT and SKIP for pagination
                result = session.run(
                    """
                    MATCH (t:TestClazz) 
                    RETURN id(t) AS node_id, t
                    ORDER BY t.name ASC
                    SKIP $offset LIMIT $batch_size
                    """,
                    offset=offset, batch_size=batch_size
                )
                result = list(result)

                # Convert results to a list
                test_classes = list(result)

                # If no more results, stop iteration
                if not test_classes:
                    break

                # Yield each test class one by one
                for record in test_classes:
                    yield _convert_node(record["t"])

                # Move to the next batch
                offset += batch_size

    def update_class(self, node_id, field_value):
        """
        Update the summarization field of the specified Clazz node.
        :param node_id:
        :param field_value:
        :return:
        """
        with self.driver.session() as session:
            session.run(
                """
                MATCH (c:Clazz)
                WHERE id(c) = $node_id
                SET c.summarization = $field_value
                """,
                node_id=node_id,
                field_value=field_value
            )
        return

    def update_method(self, node_id, field_value):
        """
        Update the summarization field of the specified Method node.
        :param node_id:
        :param field_value:
        :return:
        """
        with self.driver.session() as session:
            session.run(
                """
                MATCH (m:Method)
                WHERE id(m) = $node_id
                SET m.summarization = $field_value
                """,
                node_id=node_id,
                field_value=field_value
            )
        return

    def update_variable(self, node_id, field_value):
        """
        Update the summarization field of the specified Variable node.
        :param node_id:
        :param field_value:
        :return:
        """
        with self.driver.session() as session:
            session.run(
                """
                MATCH (v:Variable)
                WHERE id(v) = $node_id
                SET v.summarization = $field_value
                """,
                node_id=node_id,
                field_value=field_value
            )
        return

    def update_test_class(self, name, tested_clazz_fqname, tested_method_fqname, method_signature, test_report,
                          coverage_rate, coverage_lines, mutation_score, mutants, find_bug, requirement):
        """
        Update the specified TestClazz node, including tested_clazz_FQname, test_report, coverage_rate,
        coverage_lines, mutation_score, mutants fields.
        """
        with self.driver.session() as session:
            session.run(
                """
                MATCH (t:TestClazz)
                WHERE t.name = $name
                SET t.tested_clazz_FQname = $tested_clazz_fqname,
                    t.tested_method_FQname = $tested_method_fqname,
                    t.method_signature = $method_signature,
                    t.test_report = $test_report,
                    t.coverage_rate = $coverage_rate,
                    t.coverage_lines = $coverage_lines,
                    t.mutation_score = $mutation_score,
                    t.mutants = $mutants,
                    t.find_bug = $find_bug,
                    t.requirement = $requirement
                """,
                name=name,
                tested_clazz_fqname=tested_clazz_fqname,
                tested_method_fqname=tested_method_fqname,
                method_signature=method_signature,
                test_report=test_report,
                coverage_rate=coverage_rate,
                coverage_lines=coverage_lines,
                mutation_score=mutation_score,
                mutants=mutants,
                find_bug=find_bug,
                requirement=requirement
            )
        return

    def load_current_method(self):
        """
        Load the Method node corresponding to the current focal_method_id.
        :return:
        """
        with self.driver.session() as session:
            result = session.run(
                """
                MATCH (m:Method)
                WHERE id(m) = $focal_method_id
                RETURN m
                """,
                focal_method_id=self.focal_method_id
            )
            result = list(result)
            return _convert_node(result[0]["m"]) if result else None

    def search_similarity_test_class(self, test_class_name: str) -> List[tuple]:
        """
        Search for similar test classes based on Jaccard similarity.
        """

        def jaccard_similarity(set1, set2):
            intersection = len(set1.intersection(set2))
            union = len(set1.union(set2))
            return intersection / union if union != 0 else 0

        with self.driver.session() as session:
            # Get the related nodes set of the specified test class
            result = session.run(
                """
                MATCH (t:TestClazz {name: $test_class_name}) 
                RETURN t.related_nodes AS related_nodes
                """,
                test_class_name=test_class_name
            )
            record = result.single()
            if not record:
                return None

            target_related_nodes = set(record["related_nodes"])

            # Get the related nodes set of all other test classes
            result = session.run(
                """
                MATCH (t:TestClazz)
                RETURN t AS test_class, t.related_nodes AS related_nodes
                """
            )

            similarities = []
            for record in result:
                test_node = _convert_node(record["test_class"])
                if test_node.name == test_class_name:
                    continue
                assert isinstance(test_node, TestClazz)
                related_nodes = set(test_node.related_nodes)
                similarity = jaccard_similarity(target_related_nodes, related_nodes)
                if similarity < 0.1:
                    continue
                similarities.append((test_node, similarity))

            # Sort by similarity
            similarities.sort(key=lambda x: x[1], reverse=True)
            return similarities

    def search_constructor_in_clazz(self, clazz_name: str) -> List[Any]:
        """
        Query the constructor in the specified class.
        :param clazz_name:
        :return:
        """
        with self.driver.session() as session:
            result = session.run(
                """
                MATCH (n:Method {type:"CONSTRUCTOR"})<-[:CONTAIN]-(c:Clazz {name: $clazz_name})
                RETURN n
                """,
                clazz_name=clazz_name
            )
            result = list(result)
            return [_convert_node(record["n"]) for record in result]

    def method_calls_query(self, method_name: str, full_qualified_name: str = None, method_params: List[str] = None) -> \
            Union[str, List[Any]]:
        """
        Query the Method nodes called by the specified Method node.

        :param method_name: name attribute of the Method node
        :param full_qualified_name: full_qualified_name attribute of the Method node
        :param method_params: params attribute of the Method node (list)
        :return: 1. "NO MATCH FOR SOURCE METHOD" : source Method node not found
                2. "MULTIPLE MATCHES FOR SOURCE METHOD" : multiple source Method nodes found
                3. "NO MATCH FOR TARGET METHOD" : target Method node not found
                4. List : list of called Method nodes
        """
        # Step 1: Query Method node
        method_nodes = self.search_method_query(method_name, full_qualified_name, method_params)

        # Get unique Method node
        method_node = method_nodes[0]

        # Step 2: Use unique Method node ID to query INVOKE outgoing edges
        with self.driver.session() as session:
            result = session.run(
                """
                MATCH (m:Method)-[:INVOKE]->(target)
                WHERE id(m) = $method_id
                RETURN target
                """,
                method_id=int(method_node.id)
            )
            result = list(result)

        tuples = [_convert_node(record["target"]) for record in result]

        # Verify if related target Method nodes are found
        if not tuples:
            LoggerManager().logger.warning(f"No valid InvokeLine and Method pairs found for Method ID: {method_node.id}.")
            return "NO MATCH FOR TARGET METHOD"
        return tuples

    def method_usages_query(self, method_name: str, full_qualified_name: str, method_params: List[str]) -> Union[str, List[Any]]:
        """
        Query all nodes that call the specified method node and return the related list.

        :param method_name: method name
        :param full_qualified_name: fully qualified name of the method
        :param method_params: method parameter list
        :return: 1. "NO MATCH FOR SOURCE METHOD" : source Method node not found
                2. "MULTIPLE MATCHES FOR SOURCE METHOD" : multiple source Method nodes found
                3. "NO MATCH FOR TARGET METHOD" : target Method node not found
                4. List : list of calling nodes, including Method, Variable nodes
        """
        # Step 1: Uniquely determine Method node and verify
        method_nodes = self.search_method_query(method_name, full_qualified_name, method_params)

        # Get unique Method node
        method_node = method_nodes[0]

        # Step 2: Query all nodes that point to this method through INVOKE
        with self.driver.session() as session:
            result = session.run(
                """
                MATCH (target)-[:INVOKE]->(m:Method)
                WHERE id(m) = $method_id
                RETURN target
                """,
                method_id=int(method_node.id)
            )
            result = list(result)
        # Filter out Method nodes with type value "TEST_METHOD"
        filtered_tuples = []
        for record in result:
            node = _convert_node(record["target"])
            if Limit_Retrieve_Test_Case and isinstance(node, Method) and getattr(node, "type", None) == "TEST_METHOD":
                continue
            filtered_tuples.append(node)

        tuples = filtered_tuples

        # Verify if related target Method nodes are found
        if not tuples:
            LoggerManager().logger.warning(f"No valid InvokeLine and Method pairs found for Method ID: {method_node.id}.")
            return "NO MATCH FOR TARGET METHOD"
        return tuples

    def search_clazz_query(self, clazz_name: str, full_qualified_name: str = None) -> Union[str, Clazz]:
        """
        Query the specified Clazz node.

        :param full_qualified_name: full_qualified_name attribute of the Clazz node
        :param clazz_name: name attribute of the Clazz node
        :return: 1. "NO MATCH FOR SOURCE CLAZZ" : source Clazz node not found
                2. "MULTIPLE MATCHES FOR SOURCE CLAZZ" : multiple source Clazz nodes found
                3. Clazz : queried Clazz node
        """

        if not full_qualified_name:
            query = f"""
            MATCH (node:Clazz)
            WHERE node.name = $clazz_name
            RETURN node
            """
        else:
            query = f"""
            MATCH (node:Clazz)
            WHERE node.name = $clazz_name
            AND node.full_qualified_name contains $full_qualified_name
            RETURN node
            """
        if '.' in clazz_name:
            clazz_name = clazz_name.split('.')[-1]
        with self.driver.session() as session:
            result = session.run(query, {"clazz_name": clazz_name, "full_qualified_name": full_qualified_name})
            result = list(result)

        # if not result:
        #     LoggerManager().logger.error(f"No nodes found for the given clazz name: {clazz_name}.")
        #     return "NO MATCH FOR SOURCE CLAZZ"

        if len(result) > 1:
            LoggerManager().logger.error(f"Multiple nodes found for the given clazz name: {clazz_name}.")
            return "MULTIPLE MATCHES FOR SOURCE CLAZZ"

        return _convert_node(result[0]["node"])

    def search_parent_clazz_query(self, clazz_name: str) -> Union[str, List[Clazz]]:
        """
        Query the parent Clazz nodes of the specified Clazz node. (including EXTEND and IMPLEMENT relationships)
        :param clazz_name:
        :return: 1. "NO MATCH FOR PARENT CLAZZ" : parent Clazz node not found
                2. List : list of parent Clazz nodes
        """
        with self.driver.session() as session:
            result = session.run(
                """
                MATCH (n:Clazz {name: $name})-[:EXTEND|IMPLEMENT]->(m:Clazz)
                RETURN m
                """,
                name=clazz_name
            )
            result = list(result)
        # if not result:
        #     LoggerManager().logger.error(f"No parent nodes found for the given clazz name: {clazz_name}.")
        #     return "NO MATCH FOR PARENT CLAZZ"

        return [_convert_node(record["m"]) for record in result]

    def search_sub_clazz_query(self, clazz_name: str) -> Union[str, List[Clazz]]:
        """
        Query the sub Clazz nodes of the specified Clazz node. (including EXTEND and IMPLEMENT relationships)
        :param clazz_name:
        :return: 1. "NO MATCH FOR SUB CLAZZ" : sub Clazz node not found
                2. List : list of sub Clazz nodes
        """
        with self.driver.session() as session:
            result = session.run(
                """
                MATCH (n:Clazz {name: $name})<-[:EXTEND|IMPLEMENT]-(m:Clazz)
                RETURN m
                """,
                name=clazz_name
            )
            result = list(result)

        # if not result:
        #     LoggerManager().logger.error(f"No sub nodes found for the given clazz name: {clazz_name}.")
        #     return "NO MATCH FOR SUB CLAZZ"

        return [_convert_node(record["m"]) for record in result]

    def search_method_by_id(self, method_id: int) -> Union[str, Method]:
        """
        Query the specified Method node.
        :param method_id:
        :return: 1. "NO MATCH FOR SOURCE METHOD" : source Method node not found
                2. Method : queried Method node
        """
        with self.driver.session() as session:
            result = session.run(
                """
                MATCH (m:Method)
                WHERE id(m) = $method_id
                RETURN m
                """,
                method_id=method_id
            )
            result = list(result)
        # if not result:
        #     LoggerManager().logger.error(f"No nodes found for the given method ID: {method_id}.")
        #     return "NO MATCH FOR SOURCE METHOD"

        return _convert_node(result[0]["m"])

    def search_method_by_signature(self, method_signature: str) -> Union[str, Method]:
        """
        Query the specified Method node.
        :param method_signature:
        :return: 1. "NO MATCH FOR SOURCE METHOD" : source Method node not found
                2. Method : queried Method node
        """
        with self.driver.session() as session:
            result = session.run(
                """
                MATCH (m:Method)
                WHERE m.signature = $signature
                RETURN m
                """,
                signature=method_signature
            )
            result = list(result)
        # if not result:
        #     LoggerManager().logger.error(f"No nodes found for the given method signature: {method_signature}.")
        #     return "NO MATCH FOR SOURCE METHOD"

        return _convert_node(result[0]["m"])

    def search_test_class_by_signature(self, method_signature: str, fqname: str) -> Union[str, List[TestClazz]]:
        """
        Query test class nodes that test the method based on method signature.
        :param method_signature:
        :return: 1. "NO MATCH FOR SOURCE TEST CLASS" : source TestClazz node not found
                2. TestClazz : queried TestClazz node
        """
        with self.driver.session() as session:
            result = session.run(
                """
                MATCH (t:TestClazz)
                WHERE t.method_signature = $signature AND t.tested_method_FQname = $fqname
                RETURN t
                """,
                signature=method_signature,
                fqname=fqname
            )
            result = list(result)
        if not result:
            LoggerManager().logger.error(f"No Test Clazz nodes found for testing the given method signature: {method_signature}.")
            return "NO MATCH FOR SOURCE TEST CLASS"

    def delete_test_class_by_signature(self, method_signature: str, fqname: str) -> dict:
        """
        Delete test class nodes that test the method based on method signature and return detailed deletion information.

        :param method_signature: signature of the tested method
        :return: dictionary containing operation success status, number of deleted test classes, number of deleted relationships, etc.
        """
        with self.driver.session() as session:
            result = session.run(
                """
                MATCH (t:TestClazz)
                WHERE t.method_signature = $signature AND t.tested_method_FQname = $fqname
                OPTIONAL MATCH (t)-[r]-()
                DELETE r, t
                RETURN count(t) AS deleted_nodes, count(r) AS deleted_rels
                """,
                signature=method_signature, fqname=fqname
            )
            record = result.single()
            deleted_nodes = record["deleted_nodes"]
            deleted_rels = record["deleted_relationships"]

            return {
                "found": deleted_nodes > 0,
                "deleted_nodes": deleted_nodes,
                "deleted_relationships": deleted_rels,
                "success": deleted_nodes > 0
            }

    def search_method_query(self, method_name: str, full_qualified_name: str = None, method_params: List[str] = None) \
            -> Union[str, List[Method]]:
        """
        Query the specified Method node.
        :param method_name:
        :param full_qualified_name:
        :param method_params:
        :return: 1. "NO MATCH FOR SOURCE METHOD" : source Method node not found
                2. List : list of queried Method nodes
        """

        # flag = 0 means the first query only used method_name parameter
        # = 1 means used method_name and full_qualified_name parameters
        # = 2 means used method_name and method_params parameters
        # = 3 means used method_name, full_qualified_name and method_params parameters
        flag = 0
        # First query: check if nodes can be found directly
        with self.driver.session() as session:
            # Basic query part
            query = "MATCH (m:Method)"
            # Dynamically add conditions
            conditions = []
            if '.' not in method_name:
                conditions.append(f"m.name = '{method_name}'")
            else:
                conditions.append(f"m.full_qualified_name CONTAINS '{method_name}'")
            if full_qualified_name:
                conditions.append(f"m.full_qualified_name CONTAINS '{full_qualified_name}'")
                flag += 1
            if method_params:
                conditions.append(f"m.params = {method_params}")
                flag += 2
            # Concatenate conditions to query statement
            if conditions:
                query += " WHERE " + " AND ".join(conditions)
            # Return final query
            query += " RETURN m"

            result = session.run(query)
            result = list(result)
        if len(result) == 1:
            return [_convert_node(record["m"]) for record in result]
        if len(result) > 1:
            node_ids = [record["m"].id for record in result]
            if self.focal_method_id in node_ids:
                node_ids.remove(self.focal_method_id)
            with self.driver.session() as session:
                sub_result = session.run(
                    """
                    UNWIND $node_ids AS node_id
                    MATCH (base:Method), (m)
                    WHERE id(m) = node_id AND id(base) = $base_method_id
                    MATCH path = shortestPath((base)-[*..8]-(m))
                    RETURN m, length(path) AS distance
                    ORDER BY distance
                    """,
                    node_ids=node_ids,
                    base_method_id=self.focal_method_id
                )
                sub_result = list(sub_result)
            return [_convert_node(record["m"]) for record in sub_result]

        # if len(result) == 0 and flag == 0:
        #     LoggerManager().logger.error(f"No nodes found for the given method name: {method_name}.")
        #     return "NO MATCH FOR SOURCE METHOD"

        if len(result) == 0 and flag > 0:
            with self.driver.session() as session:
                result = session.run(
                    """
                    MATCH (m:Method)
                    WHERE m.name = $name
                    RETURN m
                    """,
                    name=method_name
                )
                result = list(result)
            if len(result) == 1:
                return [_convert_node(record["m"]) for record in result]
            # if len(result) == 0:
            #     LoggerManager().logger.error(f"No nodes found for the given method name: {method_name}.")
            #     return "NO MATCH FOR SOURCE METHOD"
            if len(result) > 1 and flag % 2 != 0:
                with self.driver.session() as session:
                    sub_result = session.run(
                        """
                        MATCH (m:Method)
                        WHERE m.name = $name AND m.full_qualified_name CONTAINS $full_qualified_name
                        RETURN m
                        """,
                        name=method_name, full_qualified_name=full_qualified_name
                    )
                    sub_result = list(sub_result)
                if len(sub_result) > 0:
                    sub_node_ids = [record["m"].id for record in sub_result]
                    if self.focal_method_id in sub_node_ids:
                        sub_node_ids.remove(self.focal_method_id)
                    with self.driver.session() as session:
                        sub_sub_result = session.run(
                            """
                            UNWIND $node_ids AS node_id
                            MATCH (base:Method), (m)
                            WHERE id(m) = node_id AND id(base) = $base_method_id
                            MATCH path = shortestPath((base)-[*..8]-(m))
                            RETURN m, length(path) AS distance
                            ORDER BY distance
                            """,
                            node_ids=sub_node_ids,
                            base_method_id=self.focal_method_id
                        )
                        sub_sub_result = list(sub_sub_result)
                    return [_convert_node(record["m"]) for record in sub_sub_result]
                if len(sub_result) == 0:
                    node_ids = [record["m"].id for record in result]
                    if self.focal_method_id in node_ids:
                        node_ids.remove(self.focal_method_id)
                    with self.driver.session() as session:
                        sub_result = session.run(
                            """
                            UNWIND $node_ids AS node_id
                            MATCH (base:Method), (m)
                            WHERE id(m) = node_id AND id(base) = $base_method_id
                            MATCH path = shortestPath((base)-[*..8]-(m))
                            RETURN m, length(path) AS distance
                            ORDER BY distance
                            """,
                            node_ids=node_ids,
                            base_method_id=self.focal_method_id
                        )
                        sub_result = list(sub_result)
                    return [_convert_node(record["m"]) for record in sub_result]

    def search_variable_query(self, variable_name: str, full_qualified_name: str = None) -> Union[str, List[Any]]:
        """
        Query the corresponding Variable node.

        :param full_qualified_name:
        :param variable_name: name attribute of the Variable node
        :return: 1. "NO MATCH FOR SOURCE VARIABLE" : source Variable node not found
                2. List : list of queried Variable nodes
        """
        flag = 0
        with self.driver.session() as session:
            query = "MATCH (n:Variable)"
            # Dynamically add conditions
            conditions = []
            if '.' not in variable_name:
                conditions.append(f"n.name = '{variable_name}'")
            else:
                conditions.append(f"n.full_qualified_name CONTAINS '{variable_name}'")
            if full_qualified_name:
                conditions.append(f"n.full_qualified_name CONTAINS '{full_qualified_name}'")
                flag += 1
            # Concatenate conditions to query statement
            if conditions:
                query += " WHERE " + " AND ".join(conditions)
            # Return final query
            query += " RETURN n"
            result = session.run(query)
            result = list(result)

        if len(result) == 0 and flag > 0:
            with self.driver.session() as session:
                result = session.run(
                    """
                    MATCH (n:Variable)
                    WHERE n.name = $name
                    RETURN n
                    """,
                    name=variable_name
                )
                result = list(result)
            if len(result) == 1:
                return [_convert_node(record["n"]) for record in result]
            # if len(result) == 0:
            #     LoggerManager().logger.error(f"No nodes found for the given variable name: {variable_name}.")
            #     return "NO MATCH FOR SOURCE VARIABLE"
            if len(result) > 1:
                node_ids = [record["n"].id for record in result]
                with self.driver.session() as session:
                    sub_result = session.run(
                        """
                        UNWIND $node_ids AS node_id
                        MATCH (base:Method), (n)
                        WHERE id(n) = node_id AND id(base) = $base_method_id
                        MATCH path = shortestPath((base)-[*..8]-(n))
                        RETURN n, length(path) AS distance
                        ORDER BY distance
                        """,
                        node_ids=node_ids,
                        base_method_id=self.focal_method_id
                    )
                    sub_result = list(sub_result)
                return [_convert_node(record["n"]) for record in sub_result]
        # if len(result) == 0 and flag == 0:
        #     LoggerManager().logger.error(f"No nodes found for the given variable name: {variable_name}.")
        #     return "NO MATCH FOR SOURCE VARIABLE"
        if len(result) == 1:
            return [_convert_node(record["n"]) for record in result]
        if len(result) > 1:
            node_ids = [record["n"].id for record in result]
            with self.driver.session() as session:
                sub_result = session.run(
                    """
                    UNWIND $node_ids AS node_id
                    MATCH (base:Method), (n)
                    WHERE id(n) = node_id AND id(base) = $base_method_id
                    MATCH path = shortestPath((base)-[*..8]-(n))
                    RETURN n, length(path) AS distance
                    ORDER BY distance
                    """,
                    node_ids=node_ids,
                    base_method_id=self.focal_method_id
                )
                sub_result = list(sub_result)
            return [_convert_node(record["n"]) for record in sub_result]

    def fuzzy_search(self, name: str) -> List[Any]:
        """
        Fuzzy query for Method, Variable, Clazz nodes.
        :param name:
        :return: 1. "NO MATCH FOR TARGET NODE" : target node not found
                2. List : list of queried nodes
        """
        with self.driver.session() as session:
            query = "MATCH (n)"
            # Dynamically add conditions
            conditions = []
            if '.' not in name:
                conditions.append(f"n.name = '{name}'")
            else:
                conditions.append(f"n.full_qualified_name CONTAINS '{name}'")
            # Concatenate conditions to query statement
            if conditions:
                query += " WHERE " + " AND ".join(conditions)
            # Return final query
            query += " RETURN n"
            result = session.run(query)
            result = list(result)
        # if not result:
        #     LoggerManager().logger.error(f"No nodes found for the given name: {name}.")
        #     return "NO MATCH FOR TARGET NODE"
        return [_convert_node(record["n"]) for record in result]

    def search_one_hop_neighbors(self, method_id) -> Dict[str, List[Any]]:
        """
        Get nodes adjacent to the focal_method_id node, further filter InvokeLine nodes and their connected second nodes.
        Return node lists categorized by labels (Clazz, Variable, Method).

        :return: Dictionary containing node lists for Clazz, Variable, Method labels
        """
        query = """
                MATCH (m:Method)
                WHERE id(m) = $focal_method_id

                // Part 1: Find Clazz nodes directly connected to Method
                OPTIONAL MATCH (m)-[:CONTAIN]-(n:Clazz)
                WITH COLLECT(DISTINCT n) AS direct_clazz_nodes, m

                // Part 2: Find Clazz nodes connected to Method (through DEPENDENCY relationship)
                OPTIONAL MATCH (m)-[:DEPENDENCY]-(n:Clazz)
               WITH direct_clazz_nodes, COLLECT(DISTINCT n) AS related_clazz_nodes, m

                // Find InvokeLine nodes and get other nodes connected to them
                OPTIONAL MATCH (m)-[]->(vo:Variable)
                WHERE NOT id(vo) = id(m) // Exclude returning to original Method node

                OPTIONAL MATCH (m)-[]->(mo:Method)
                WHERE NOT id(mo) = id(m) // Exclude returning to original Method node

                OPTIONAL MATCH (m)<-[]-(vi:Variable)
                WHERE NOT id(vi) = id(m) // Exclude returning to original Method node

                OPTIONAL MATCH (m)<-[]-(mi:Method)
                WHERE NOT id(mi) = id(m) // Exclude returning to original Method node
                
                // Aggregate all nodes
                WITH direct_clazz_nodes, 
                     related_clazz_nodes,
                     COLLECT(DISTINCT vo) AS variable_nodes_outgoing,
                     COLLECT(DISTINCT vi) AS variable_nodes_incoming,
                     COLLECT(DISTINCT mo) AS method_nodes_outgoing,
                     COLLECT(DISTINCT mi) AS method_nodes_incoming
                     
                // Final result: merge direct_clazz_nodes and clazz_nodes_from_invoke
                RETURN 
                    direct_clazz_nodes AS main_clazz,
                    related_clazz_nodes AS clazz_nodes,
                    variable_nodes_outgoing,
                    variable_nodes_incoming,
                    method_nodes_outgoing,
                    method_nodes_incoming
                """

        parameters = {"focal_method_id": method_id}

        try:
            with self.driver.session() as session:
                result = session.run(query, parameters)
                record = result.single()
                # Convert Neo4j nodes to corresponding class instances
                main_clazz = [_convert_node(node) for node in record["main_clazz"] or []]
                clazz_nodes = [_convert_node(node) for node in record["clazz_nodes"] or []]
                variable_nodes_outgoing = [_convert_node(node) for node in
                                           record["variable_nodes_outgoing"] or []]
                variable_nodes_incoming = [_convert_node(node) for node in
                                           record["variable_nodes_incoming"] or []]
                method_nodes_outgoing = [_convert_node(node) for node in record["method_nodes_outgoing"] or []]
                method_nodes_incoming = [_convert_node(node) for node in record["method_nodes_incoming"] or []]

                return {
                    "MainClazz": main_clazz,
                    "Clazz": clazz_nodes,
                    "VariableOutgoing": variable_nodes_outgoing,
                    "VariableIncoming": variable_nodes_incoming,
                    "MethodOutgoing": method_nodes_outgoing,
                    "MethodIncoming": method_nodes_incoming
                }
        except Neo4jError as e:
            LoggerManager().logger.error(f"Error while querying adjacent nodes: {e}")
            LoggerManager().logger.exception("Exception occurred")
            return {}

    def search_variable_one_hop_neighbors(self, field_id) -> Dict[str, List[Any]]:
        """
        Get nodes adjacent to the focal_method_id node, further filter InvokeLine nodes and their connected second nodes.
        Return node lists categorized by labels (Clazz, Variable, Method).

        :return: Dictionary containing node lists for Clazz, Variable, Method labels
        """
        query = """
                MATCH (m:Variable)
                WHERE id(m) = $focal_variable_id

                // Part 1: Find Clazz nodes directly connected to Variable
                OPTIONAL MATCH (m)-[:CONTAIN]-(n:Clazz)
                WITH COLLECT(DISTINCT n) AS direct_clazz_nodes, m

                // Part 2: Find Clazz nodes connected to Variable (through DEPENDENCY relationship)
                OPTIONAL MATCH (m)-[:DEPENDENCY]-(n:Clazz)
               WITH direct_clazz_nodes, COLLECT(DISTINCT n) AS related_clazz_nodes, m

                // Find InvokeLine nodes and get other nodes connected to them
                OPTIONAL MATCH (m)-[]->(vo:Variable)
                WHERE NOT id(vo) = id(m) // Exclude returning to original Method node

                OPTIONAL MATCH (m)-[]->(mo:Method)
                WHERE NOT id(mo) = id(m) // Exclude returning to original Method node

                OPTIONAL MATCH (m)<-[]-(vi:Variable)
                WHERE NOT id(vi) = id(m) // Exclude returning to original Method node

                OPTIONAL MATCH (m)<-[]-(mi:Method)
                WHERE NOT id(mi) = id(m) // Exclude returning to original Method node
                
                // Aggregate all nodes
                WITH direct_clazz_nodes, 
                     related_clazz_nodes,
                     COLLECT(DISTINCT vo) AS variable_nodes_outgoing,
                     COLLECT(DISTINCT vi) AS variable_nodes_incoming,
                     COLLECT(DISTINCT mo) AS method_nodes_outgoing,
                     COLLECT(DISTINCT mi) AS method_nodes_incoming
                     
                // Final result: merge direct_clazz_nodes and clazz_nodes_from_invoke
                RETURN 
                    direct_clazz_nodes AS main_clazz,
                    related_clazz_nodes AS clazz_nodes,
                    variable_nodes_outgoing,
                    variable_nodes_incoming,
                    method_nodes_outgoing,
                    method_nodes_incoming
                """

        parameters = {"focal_variable_id": field_id}

        try:
            with self.driver.session() as session:
                result = session.run(query, parameters)
                record = result.single()
                # Convert Neo4j nodes to corresponding class instances
                main_clazz = [_convert_node(node) for node in record["main_clazz"] or []]
                clazz_nodes = [_convert_node(node) for node in record["clazz_nodes"] or []]
                variable_nodes_outgoing = [_convert_node(node) for node in
                                           record["variable_nodes_outgoing"] or []]
                variable_nodes_incoming = [_convert_node(node) for node in
                                           record["variable_nodes_incoming"] or []]
                method_nodes_outgoing = [_convert_node(node) for node in record["method_nodes_outgoing"] or []]
                method_nodes_incoming = [_convert_node(node) for node in record["method_nodes_incoming"] or []]

                return {
                    "MainClazz": main_clazz,
                    "Clazz": clazz_nodes,
                    "VariableOutgoing": variable_nodes_outgoing,
                    "VariableIncoming": variable_nodes_incoming,
                    "MethodOutgoing": method_nodes_outgoing,
                    "MethodIncoming": method_nodes_incoming
                }
        except Neo4jError as e:
            LoggerManager().logger.error(f"Error while querying adjacent nodes: {e}")
            LoggerManager().logger.exception("Exception occurred")
            return {"Clazz": [], "Variable": [], "Method": []}

    # def handle_1_self_graph(self, method_dict):
    #     """
    #     Handle information in the 1-self graph, including Method, Variable, Clazz nodes. [Not used currently]
    #     :param method_dict:
    #     :return:
    #     """
    #     related_class = method_dict['Clazz']
    #     related_method = method_dict['Method']
    #     related_variable = method_dict['Variable']
    #     current_method = self.load_current_method()
    #     related_method.append(current_method)
    #     used_list = []
    #     main_class = None
    #     surround_class = []
    #     for clazz in related_class:
    #         content = clazz.simple_content
    #         constructor_list = self.search_constructor_in_clazz(clazz.name)
    #         for constructor in constructor_list:
    #             content = self.append_info_to_class(content, constructor.content)
    #         for variable in related_variable:
    #             if variable.class_name == clazz.name:
    #                 content = self.append_info_to_class(content, variable.content)
    #                 used_list.append(variable)
    #         for method in related_method:
    #             if method.class_name == clazz.name:
    #                 content = self.append_info_to_class(content, method.content)
    #                 used_list.append(method)
    #         if clazz.name == current_method.class_name:
    #             main_class = content
    #         else:
    #             surround_class.append(content)
    #
    #     remain_method = [self.construct_info_in_class(item.class_name, item.content) for item in related_method if
    #                      item not in used_list]
    #     remain_variable = [self.construct_info_in_class(item.class_name, item.content) for item in related_variable if
    #                        item not in used_list]
    #
    #     return main_class, surround_class, remain_method, remain_variable
    #
